from Crypto import Random
from Crypto.Signature import pkcs1_15
from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA

from Crypto.Signature import PKCS1_v1_5 as PKCS1_signature  # 检测
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher  # 加密解密


import base64

class RsaUtil:

    private_key_pem = """-----BEGIN RSA PRIVATE KEY-----
        MIIEogIBAAKCAQEAnJgNow1ws5u/wXLKjWZrkM5bIEq+SNgxz/FbeDzpMCQAbDMF
        3A5hCJKMLyyadD3Brez6hUnGpC6ysRS4WeIDIfl1A3ZhPThlazQbxzcw+lThBpGV
        a5wnbk2/E+L46SCImJh2PGuQGTn6qfUnJ1PQS3L/B6KXU1au41XymMIkUgKa68xP
        QYspoNCeSRGmy2yO+kt9FXsdSXmdar/LvATsJQvuSueKm43ZPfebdMrTJGM/oJjX
        ywsHRUwZuE1MKFOijvu9oSIDdYPTeejveUZQ60zgpbt0bANccTLQZdoPPh27QBuc
        4A7thCLLusamNd4ioItClV9Dn0Oa66I/aS1u+wIDAQABAoIBAC2IlzqIz3MigD0v
        dE9ubyuxnVTD9xW2a7GT9G1/hKzfYxITOgp1UejMuB+FgDn5tSYuSOBWBgE6pPIR
        3NMiAZgw8cH6e6byDaeBBOzGEyNPaAlhbsnWI68Sw1N0hb1AHQgGGSdTfZigP8+I
        vgRU/jXaL+cX85mjlIH9PWgG18nSjI4o+Y/SXEKGhhDDGt1teMs/p069bWZ4T0TW
        P1OSZm8ZiyHStuKTHNSTjdGlknQgcOogxGfchh9qAI112JA9v1AD1kAtatWDT1mk
        R5UZS6tESsm3WoATo1mk6XOytqnyH/xhi7FnVhykzNoRg+3I0FuxyUz/c+JHixbS
        Z+sxlcUCgYEAv7p+vagE0qBCk3eZFo2LGCVeuuBJ9HmCYW5taxh2L86fiuxRGRdp
        PZm0HReRZzvmRKBuz73v0FGGUXQAOQi/gmM7jzsm7QbaGcaT/GYSmJbzPavnz9fU
        NIexnfOws9g96J6tadZpK7+blFJ3K/+x43VIdj/J5g5pmjhvnxHcAW0CgYEA0RZt
        rCJSUQ10Qv6J8esA47IeZUElGBNHhkkhA1ZZ0DujBtpm2gNqDFkhXCHeac9Qlo/4
        SucHgfCwGA6Wb57D2uIEg90VAgOROxaUGu891rj0twaPXO5SnYAj7pqmC0pZ9pX9
        wFPjfc6CuRCfn4NYrUq79r2k/ZJq5pmc/4ZE2QcCgYBUYhReDJDy0q36kWzDTFCn
        s+rZfNWPz9reXuexdvatW9dEIXYgAPnErt1i4HiSYv6ZvbKHpCTkHr1rJA7lC0Ce
        dmZER9h+AcIEDi4brCELH4S24+4GwAsZg49c4WAoB0hYFV5lcoSfSJmE0ja5GrSE
        U54szS5jHnzD74rRC980kQKBgH8AoJ1uO5IAGh6XfjL0DBZCmzPl3iJBk7B9uw/L
        48IYd24yO0cxy7i0aw8Lg43B37GUwby9mFHVYSE7uyURyuWpZBS0tY4DnLPHgy7Z
        IzbKIdsKoQpRuNzyeOD2DuXMw9iEQx4P1p8UUVMYRd0eaz7XplS53P4pccWgiTys
        0iUBAoGAXn0ajgHEONi11DwoJ0moNvcnNE0vnGjjr07LaMXiDeJYKZyiaaWt6J7E
        J9s2cxZ0QdE3Z2sKwEBrkwHf13zYgPZFWlgA3O1QjhVkhXghmCpuHalmj7xQizDj
        ZXgyGGs/nsb0GQ654CSFEtrXrWg2eSEenZcUnV8ti4WtufgBp+k=
        -----END RSA PRIVATE KEY-----"""

    public_key_pem = """-----BEGIN PUBLIC KEY-----
        MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAnJgNow1ws5u/wXLKjWZr
        kM5bIEq+SNgxz/FbeDzpMCQAbDMF3A5hCJKMLyyadD3Brez6hUnGpC6ysRS4WeID
        Ifl1A3ZhPThlazQbxzcw+lThBpGVa5wnbk2/E+L46SCImJh2PGuQGTn6qfUnJ1PQ
        S3L/B6KXU1au41XymMIkUgKa68xPQYspoNCeSRGmy2yO+kt9FXsdSXmdar/LvATs
        JQvuSueKm43ZPfebdMrTJGM/oJjXywsHRUwZuE1MKFOijvu9oSIDdYPTeejveUZQ
        60zgpbt0bANccTLQZdoPPh27QBuc4A7thCLLusamNd4ioItClV9Dn0Oa66I/aS1u
        +wIDAQAB
        -----END PUBLIC KEY-----"""

    @staticmethod
    def generate_key():
        random_generator = Random.new().read
        rsa = RSA.generate(2048, random_generator)
        # 生成私钥
        private_key = rsa.exportKey()
        # 生成公钥
        public_key = rsa.publickey().exportKey()
        return private_key.decode('utf-8'), public_key.decode('utf-8')

    @staticmethod
    def load_public_key(pem_data):
        """从 PEM 字符串加载公钥"""
        return RSA.import_key(pem_data)

    @staticmethod
    def load_private_key(pem_data):
        """从 PEM 字符串加载私钥"""
        return RSA.import_key(pem_data)

    @staticmethod
    def sign(data: str, private_key_pem= private_key_pem) -> str:
        """
        使用私钥对数据签名
        :param private_key_pem: 私钥 PEM 字符串
        :param data: 要签名的原始字符串
        :return: Base64 编码的签名值
        """
        key = RsaUtil.load_private_key(private_key_pem)
        h = SHA256.new(data.encode("utf-8"))
        signer = pkcs1_15.new(key)
        signature = signer.sign(h)
        return base64.b64encode(signature).decode("utf-8")

    @staticmethod
    def verify(data: str, signature_b64: str, public_key_pem=public_key_pem) -> bool:
        """
        验证签名是否有效
        :param public_key_pem: 公钥 PEM 字符串
        :param data: 原始数据
        :param signature_b64: Base64 编码的签名值
        :return: True if valid, False otherwise
        """
        key = RsaUtil.load_public_key(public_key_pem)
        h = SHA256.new(data.encode("utf-8"))
        verifier = pkcs1_15.new(key)
        try:
            verifier.verify(h, base64.b64decode(signature_b64))
            return True
        except (ValueError, TypeError):
            return False

    @staticmethod
    def rsa_public_sign(data, _public_key=public_key_pem):
        public_key = RsaUtil.load_public_key(_public_key)
        signer = PKCS1_cipher.new(public_key)
        _sign = signer.encrypt(data.encode("utf-8"))
        signature = base64.b64encode(_sign)
        signature = signature.decode('utf-8')
        return signature


    @staticmethod
    def rsa_private_parse(_sign, _private_key=private_key_pem):
        pri_key = RsaUtil.load_private_key(_private_key)
        cipher = PKCS1_cipher.new(pri_key)
        back_text = cipher.decrypt(base64.b64decode(_sign), "123", 0)
        return back_text.decode('utf-8')


if __name__ == "__main__":
    # print(g__private_key)
    # print(g__public_key)
    data = "admin123"

    signature = RsaUtil.sign(data)
    print("Signature:", signature)

    is_valid = RsaUtil.verify(data, signature)
    print("Signature valid?", is_valid)
